
import numpy as np

class KNearest:
    def __init__(self):
        pass
    
    def train(self, dataSet, labels, clss=None):
        self.dataSet = dataSet
        self.labels = labels
        if clss==None:
            clss = labels
        self.clss = clss
        self.nclss = labels.max()+1
        self.init()

    def init(self):
        aa = self.labels.astype('int32').flatten().max()
        self.max_id = aa+1

    def save(self, fn):
        np.savez(fn, dataSet=self.dataSet, labels=self.labels, clss=self.clss)

    def nclss(self):
        return self.labels.max()+1

    def load(self, fn):
        data = np.load(fn)
        self.dataSet = data['dataSet']
        self.labels = data['labels']
        self.clss = data['clss']
        self.init()

    def findNearest(self, inX, k):
        results=[]
        neighbours=[]
        dist=[]
        for i in range(inX.shape[0]):
            results1, neighbours1, dist1 = self.findNearest1(inX[i,:], k)
            results.append(results1)
            neighbours.append(neighbours1)
            dist.append(dist1)
        results = np.array(results)
        neighbours = np.mat(neighbours)
        dist = np.mat(dist)
        ret = []
        for x in results:
            ret.append(x)
        return ret, results, neighbours, dist

    def findNearest1(self, inX, k):
        dataSetSize = self.dataSet.shape[0]
        diffMat = np.tile(inX, (dataSetSize, 1)) - self.dataSet
        sqDiffMat = diffMat**2
        sqDistances = sqDiffMat.sum(axis=1)
        distances = sqDistances**0.5
        sortedDistIndicies = distances.argsort()
        classCount = np.zeros((self.max_id))
        results = -1
        neighbours = []
        dist = []
        for i in range(k):
            j = sortedDistIndicies[i]
            voteIlabel = int(self.labels[j])
            neighbours.append(voteIlabel)
            dist.append(distances[j])
            classCount[voteIlabel] += 1
        cc = classCount.argsort()
        results = cc[-1]
        return results, neighbours, dist



